import cv2
import numpy as np

cap = cv2.VideoCapture(0)
# 置信度阈值
ConThreshold = 0.5
# 非极大值抑制阈值
nmsThreshold = 0.3
whT = 320
# ImgDemo = "C:\\Users\\86191\\Pictures\\Saved Pictures\\Camera Roll\\car.jpg"
ClassDemo = 'coco.names' # 加载类别文件
ClassData = []
with open(ClassDemo,'rt') as f:
    ClassData = f.read().rstrip('\n').split('\n')

ModelCfg = 'yolov3.cfg'
ModelWeights = 'yolov3.weights'
"""
cv2.dnn.readNetFromDarknet(ModelCfg, ModelWeights) 从指定的配置文件和权重文件中读取网络模型。其中ModelCfg是模型配置文件的路径，ModelWeights是模型权重文件的路径。
net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV) 设置网络的后端为OpenCV。这指定使用OpenCV的深度神经网络模块进行网络的计算。
net.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU) 设置网络的目标为CPU。这意味着网络的计算将在CPU上进行
"""
net = cv2.dnn.readNetFromDarknet(ModelCfg, ModelWeights) # 将模型添加到dnn网络中
net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV) # 设置网络后端为opencv
net.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU) # 设置网络为CPU驱动


def findObjects(outputs,frame):
    hT,wT,cT  = frame.shape
    bbox = []
    classIds = []
    confid = []
    for output in outputs:
        for det in output:
            scores = det[5:]
            classId = np.argmax(scores)
            confidence = scores[classId]
            if confidence > ConThreshold:
                w,h = int(det[2]*wT),int(det[3]*wT)
                x,y = int((det[0]*wT)-w/2),int((det[1]*hT)-h/2)
                # 矩形框的坐标
                bbox.append([x,y,w,h])
                classIds.append(classId)
                confid.append(float(confidence))
    # 非极大值抑制
    indice = cv2.dnn.NMSBoxes(bbox,confid,ConThreshold,nmsThreshold)
    # print(indice)
    for i in indice:

        box = bbox[i]
        x,y,w,h = box[0],box[1],box[2],box[3]
        # 绘制矩形
        cv2.rectangle(frame,(x,y),(x+w,y+h),(0,255,0),2)
        cv2.putText(frame,f'{ClassData[classIds[i]]} {int(confid[i]*100)}%',(x,y-10),cv2.FONT_HERSHEY_SIMPLEX,0.6,(0,0,255),2)



while True:
    ret,frame = cap.read()
    # blob特定格式
    blob = cv2.dnn.blobFromImage(frame, 1 / 255, (whT, whT), [0, 0, 0], True, crop=False)
    # 将输出的blob作为传入网络的输入
    net.setInput(blob)
    # 获取输出层的名称
    layerNames = net.getLayerNames()
    # print(layerNames)
    # print(net.getUnconnectedOutLayers())
    # 获取输出层的最后一个
    OutputNames = [layerNames[i-1] for i in net.getUnconnectedOutLayers()]
    # print(OutputNames)

    Outputs = net.forward(OutputNames)
    # print(Outputs[0].shape)
    # print(Outputs[1].shape)
    # print(Outputs[2].shape)
    # print(Outputs[0][0])
    findObjects(Outputs,frame)
    cv2.imshow('Demo',frame)
    cv2.waitKey(1)
    if cv2.waitKey(1) & 0xFF == ord('q'):
         break

cap.release()
cv2.destroyAllWindows()